from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from typing import List
from app.database import get_db
from app import models, schemas, auth
from app.llm_service import generate_response_stream
import json

router = APIRouter(prefix="/sessions", tags=["会话"])


@router.post("", response_model=schemas.Response)
async def create_session(
    request: schemas.SessionCreate,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """创建会话"""
    # 检查Agent是否存在
    agent = db.query(models.Agent).filter(models.Agent.id == request.agent_id).first()
    if not agent:
        raise HTTPException(status_code=404, detail="Agent不存在")
    
    session = models.Session(
        user_id=current_user.id,
        agent_id=request.agent_id,
        title=request.title or "新会话"
    )
    db.add(session)
    db.commit()
    db.refresh(session)
    
    return schemas.Response(
        message="会话创建成功",
        data=schemas.SessionResponse.from_orm(session)
    )


@router.get("", response_model=schemas.Response)
async def list_sessions(
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """获取当前用户的会话列表"""
    sessions = db.query(models.Session).filter(
        models.Session.user_id == current_user.id,
        models.Session.status == "active"
    ).order_by(
        models.Session.pinned.desc(),
        models.Session.updated_at.desc()
    ).all()
    
    return schemas.Response(
        data=[schemas.SessionResponse.from_orm(s) for s in sessions]
    )


@router.get("/{session_id}", response_model=schemas.Response)
async def get_session(
    session_id: int,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """获取会话详情"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    return schemas.Response(data=schemas.SessionResponse.from_orm(session))


@router.patch("/{session_id}", response_model=schemas.Response)
async def update_session(
    session_id: int,
    request: schemas.SessionUpdate,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """更新会话"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    if request.title is not None:
        session.title = request.title
    if request.pinned is not None:
        session.pinned = request.pinned
    
    db.commit()
    db.refresh(session)
    
    return schemas.Response(
        message="更新成功",
        data=schemas.SessionResponse.from_orm(session)
    )


@router.delete("/{session_id}", response_model=schemas.Response)
async def delete_session(
    session_id: int,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """删除会话"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    session.status = "deleted"
    db.commit()
    
    return schemas.Response(message="删除成功")


@router.get("/{session_id}/messages", response_model=schemas.Response)
async def list_messages(
    session_id: int,
    page: int = Query(1, ge=1, description="页码"),
    page_size: int = Query(10, ge=1, le=100, description="每页消息数"),
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """获取会话消息列表（分页，按id倒序查询）"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    # 查询总数
    total = db.query(models.Message).filter(
        models.Message.session_id == session_id
    ).count()
    
    # 按id倒序分页查询
    messages = db.query(models.Message).filter(
        models.Message.session_id == session_id
    ).order_by(models.Message.id.desc()).offset(
        (page - 1) * page_size
    ).limit(page_size).all()
    
    # 倒转结果（前端显示时按升序）
    messages.reverse()
    
    return schemas.Response(
        data={
            "messages": [schemas.MessageResponse.from_orm(m) for m in messages],
            "total": total,
            "page": page,
            "page_size": page_size,
            "has_more": (page * page_size) < total
        }
    )


@router.post("/{session_id}/messages", response_model=schemas.Response)
async def create_message(
    session_id: int,
    request: schemas.MessageCreate,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """发送消息（非流式）"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    # 保存用户消息
    user_message = models.Message(
        session_id=session_id,
        role="user",
        content=request.content
    )
    db.add(user_message)
    db.commit()
    db.refresh(user_message)
    
    return schemas.Response(
        message="消息已发送",
        data=schemas.MessageResponse.from_orm(user_message)
    )


@router.get("/{session_id}/generate/stream")
async def generate_stream(
    session_id: int,
    prompt: str,
    current_user: models.User = Depends(auth.get_current_user),
    db: Session = Depends(get_db)
):
    """流式生成回复（SSE）"""
    session = db.query(models.Session).filter(
        models.Session.id == session_id,
        models.Session.user_id == current_user.id
    ).first()
    
    if not session:
        raise HTTPException(status_code=404, detail="会话不存在")
    
    # 获取Agent配置
    agent = session.agent
    
    # 获取历史消息
    messages = db.query(models.Message).filter(
        models.Message.session_id == session_id
    ).order_by(models.Message.created_at.asc()).all()
    
    async def event_generator():
        """SSE事件生成器"""
        try:
            import asyncio
            async for chunk in generate_response_stream(agent, messages, prompt, db, session_id):
                # 发送 chunk 事件
                event_data = f"event: chunk\ndata: {json.dumps(chunk, ensure_ascii=False)}\n\n"
                yield event_data
                # 添加小延迟确保流式效果
                await asyncio.sleep(0.01)
            
            # 发送完成事件
            yield f"event: final\ndata: {json.dumps({'status': 'completed'}, ensure_ascii=False)}\n\n"
        except Exception as e:
            error_data = {"error": str(e)}
            yield f"event: error\ndata: {json.dumps(error_data, ensure_ascii=False)}\n\n"
    
    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream; charset=utf-8",
        headers={
            "Cache-Control": "no-cache, no-transform",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
            "Content-Type": "text/event-stream; charset=utf-8"
        }
    )

